# -*- coding: utf-8 -*-
"""
mathmodeling_SAK.Plot.plot3D
============================

常用的 3D 绘图方法 
"""
# 数据分析三巨头
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from scipy.interpolate import griddata

from mpl_toolkits import mplot3d
from mpl_toolkits.mplot3d import Axes3D

plt.rcParams.update({
    "figure.dpi":100, 
    "font.sans-serif":'SimHei', 
    "axes.unicode_minus":False })

def df_grid_3D_plot(
    df, x='', y='', z='', 
    type='surface',
    figure_kwargs={},
    plot_kwargs={},
    ax_3D=None
    ):
    """ 快速绘制三维平面

    输入参数
    ----------
    df : pandas.DataFrame
        需要绘制三维视图的数据框，至少应当包含三列数据。
    x : str, optional
        需要绘制的三维数据分布情况的 x 轴数据，默认为空字符串
    y : str, optional
        需要绘制的三维数据分布情况的 y 轴数据，默认为空字符串
    z : str, optional
        需要绘制的三维数据分布情况的 z 轴数据，默认为空字符串
    type : str, optional
        绘制的图形样式种类，可以是曲面或者网格的样式。
        
        曲面样式：`surface` / `sf` / `s`
        网格样式：`wireframe` / `wf` / `w`
        
        by default 'surface'
    figure_kwargs : dict, optional
        创建的 figure 对象的参数, by default {}
    plot_kwargs : dict, optional
        绘图测参数, by default {}
    ax_3D : _type_, optional
        matplotlib Axes3D 对象, by default None
        指定绘图坐标轴，如果缺省会自动生成一个 ax。
        
    注意
    ----------
        
    你只能通过 `matplotlib.add_subplot()` 创建 Axes3D 
    对象，否则 `projection='3d'` 将不会生效。
        
    >>> import matplotlib.pyplot as plt
    >>> from mpl_toolkits.mplot3d import Axes3D
    >>> fig = plt.figure(**figure_kwargs)
    >>> ax = fig.add_subplot(111, projection='3d')
    """
    # 提取 X, Y, Z 列的值
    X = df[x].values
    Y = df[y].values
    Z = df[z].values

    # 创建网格数据
    xi = np.linspace(X.min(), X.max(), 1000)
    yi = np.linspace(Y.min(), Y.max(), 1000)
    X_grid, Y_grid = np.meshgrid(xi, yi)
    Z_grid = griddata(
        (X, Y), Z, (X_grid, Y_grid), 
        method='cubic')

    # 创建三维坐标轴对象
    if ax_3D==None:
        fig = plt.figure(**figure_kwargs)
        ax_3D = fig.add_subplot(111, projection='3d')
    
    # 绘制三维曲面图
    if type=='surface' or type=='sf' or type=='s':
        ax_3D.plot_surface(
            X_grid, Y_grid, Z_grid, **plot_kwargs)
    elif type=='wireframe' or type=='wf' or type=='w':
        ax_3D.plot_wireframe(
            X_grid, Y_grid, Z_grid, **plot_kwargs)
    else:
        print(type, ": No such type of plot\n")
    
    surf = ax_3D.plot_surface(
        X_grid, Y_grid, Z_grid, 
        **plot_kwargs)
    
def plot_sphere(
    center=(0, 1), radius=1, 
    ax=None,
    figure_kwargs={},
    plot_kwargs={},
    ):
    r""" 在三维图中绘制一个球面
    
    输入参数
    ----------
    center : tuple like, optional
        球心坐标值，形式为 :math:`\left(x, y, z\right)`, 
        by default (0, 1)
    radius : int, optional
        绘制球形的半径, by default 1
    ax : _type_, optional
        matplotlib Axes3D 对象, by default None
        指定绘图坐标轴，如果缺省会自动生成一个 ax。
    figure_kwargs : dict, optional
        创建的 figure 对象的参数, by default {}
    plot_kwargs : dict, optional
        创建的 wireframe 的参数, by default {}

    返回值
    ----------
    ax : mpl_toolkits.mplot3d.axes3d.Axes3D
        经过绘图的 3D 坐标轴
        
    注意
    ----------
        
    你只能通过 `matplotlib.add_subplot()` 创建 Axes3D 
    对象，否则 `projection='3d'` 将不会生效。
        
    >>> import matplotlib.pyplot as plt
    >>> from mpl_toolkits.mplot3d import Axes3D
    >>> fig = plt.figure(**figure_kwargs)
    >>> ax = fig.add_subplot(111, projection='3d')
    """
    
    if ax == None:
        fig = plt.figure(**figure_kwargs)
        ax = fig.add_subplot(111, projection='3d')
        
    
    # Ensure the radius is a scalar
    if isinstance(radius, (list, tuple, np.ndarray)):
        if len(radius) == 1:
            radius = radius[0]
        else:
            raise TypeError(
                "Radius must be a scalar value.")
    
    # 从数据点构建网格对象
    u = np.linspace(0, 2 * np.pi, 100)
    v = np.linspace(0, np.pi, 100)
    
    # 根据球面的参数编写方程式
    x = center[0] + (radius * np.outer(
        np.cos(u), np.sin(v)))
    y = center[1] + (radius * np.outer(
        np.sin(u), np.sin(v)))
    z = center[2] + (radius * np.outer(
        np.ones(np.size(u)), np.cos(v)))
    
    # P绘制球面
    ax.plot_wireframe(
        x, y, z, **plot_kwargs)
    
    return ax